// Copyright 2022 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only

#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include <stdio.h>
#include <pthread.h>

#include <openenclave/attestation/attester.h>
#include <openenclave/attestation/custom_claims.h>
#include <openenclave/attestation/sgx/evidence.h>

#include <noise/protocol/handshakestate.h>
#include <noise/protocol/cipherstate.h>

#include "sharded_ohtable/sharded_ohtable.h"
#include "ohtable/ohtable.h"
#include "ratelimit/ratelimit.h"
#include "aci_pni_uak_helpers/aci_pni_uak_helpers.h"
#include "noiseutil/noise.h"
#include "util/util.h"
#include "proto/cdsi.h"
#include "proto/pbutil.h"

#include "fixedset/fixedset.h"
#include "shim.h"

// Include the trusted cds header that is generated
// during the build. This file is generated by calling the
// sdk tool oeedger8r against the cds.edl file.
#include "trust/cds_t.h"

#include "enc_log.h"

void __wrap_noise_rand_bytes(void* bytes, size_t size) {
  CHECK(OE_OK == oe_random(bytes, size));
}

int getentropy(void *buf, size_t len)
{
  return oe_random(buf, len);
}

static sharded_ohtable *g_table = NULL;
static size_t g_num_shards = 0;
static size_t g_handshakestart_size = 0;
static uint8_t g_handshakestart[32 << 10] = {0};
static uint8_t g_privkey[NOISE_KEY_SIZE] = {0};
static pthread_rwlock_t g_handshakestart_mu = PTHREAD_RWLOCK_INITIALIZER;
static oe_uuid_t sgx_remote_uuid = {OE_FORMAT_UUID_SGX_ECDSA};
static uint64_t g_client_secret = 0;
static uint8_t g_init_done = 0;

#define ENCLAVE_INIT_UNINITIALIZED 0
#define ENCLAVE_INIT_INITIALIZING 1
#define ENCLAVE_INIT_INITIALIZED 2

#define CLIENT_UNUSED 0
#define CLIENT_INUSE 1
#define CLIENT_DELETE 2

static error_t check_init_complete() {
  if (__atomic_load_n(&g_init_done, __ATOMIC_ACQUIRE) != ENCLAVE_INIT_INITIALIZED) return err_ENCLAVE__GENERAL__NOINIT;
  return err_SUCCESS;
}

typedef struct
{
  // [canary] and [state] are for security, not just correctness.  Without them,
  // a host-side process might do the following:
  //
  //   cli = enclave_new_client(...)
  //   in_parallel:
  //     enclave_close_client(cli)
  //     enclave_other_client_function(cli)
  //
  // or
  //
  //   cli = enclave_new_client(...)
  //   enclave_close_client(cli)
  //   enclave_close_client(cli)  // double-free
  //
  // These could relatively easily force processing on unallocated heap space,
  // and the code is not hardened for that.  With these values, we attempt to
  // protect against that.
  uint64_t canary;
  uint8_t state;

  uint8_t *workspace;
  struct org_signal_cdsi_client_request_t *req;
  NoiseHandshakeState *handshake;
  NoiseCipherState *send;
  NoiseCipherState *recv;
} client_t;

static error_t validate_privkey(size_t key_size, uint8_t key[])
{
  if (key_size != NOISE_KEY_SIZE)
    return err_ENCLAVE__GENERAL__INVALID_PRIVATE_KEY;
  uint8_t out = 0;
  for (size_t i = 0; i < NOISE_KEY_SIZE; i++)
  {
    out |= key[i];
  }
  // private key all zeros, it hasn't been correctly initiated
  if (out == 0)
    return err_ENCLAVE__GENERAL__INVALID_PRIVATE_KEY;
  return err_SUCCESS;
}

static error_t generate_privkey(size_t key_size, uint8_t key[])
{
  ASSERT_ERR(OE_OK == oe_random(key, key_size), err_ENCLAVE__GENERAL__ENTROPY);
  key[0] &= 0xF8;
  key[31] = (key[31] & 0x7F) | 0x40;
  return validate_privkey(key_size, key);
}

// All client_ return with the client's `mu` locked.

static error_t client_setstate(client_t* c, uint8_t from_state, uint8_t to_state) {
  ASSERT_ERR(__atomic_compare_exchange_n(&c->state, &from_state, to_state, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQ_REL), err_ENCLAVE__GENERAL__CLIENT_LOCK_FAILED);
  return err_SUCCESS;
}

static error_t client_add(client_t* c, uint64_t* id) {
  ASSERT_ERR(oe_is_within_enclave(c, sizeof(client_t)), err_ENCLAVE__GENERAL__CLIENT_ADD_FAILED);
  RETURN_IF_ERROR(client_setstate(c, CLIENT_UNUSED, CLIENT_INUSE));
  *id = (uint64_t)c;
  __atomic_store_n(&c->canary, g_client_secret, __ATOMIC_RELEASE);
  TEST_LOG("client_add(%p)", c);
  return err_SUCCESS;
}

static error_t client_get(uint64_t id, client_t** c) {
  *c = (client_t*)id;
  ASSERT_ERR(oe_is_within_enclave(*c, sizeof(client_t)), err_ENCLAVE__GENERAL__CLIENT_GET_FAILED);
  ASSERT_ERR(((uintptr_t)*c) != ((uintptr_t)&g_client_secret), err_ENCLAVE__GENERAL__CLIENT_GET_FAILED);
  ASSERT_ERR(__atomic_load_n(&(*c)->canary, __ATOMIC_ACQUIRE) == g_client_secret, err_ENCLAVE__GENERAL__CLIENT_GET_FAILED);
  RETURN_IF_ERROR(client_setstate(*c, CLIENT_UNUSED, CLIENT_INUSE));
  TEST_LOG("client_get(%p)", *c);
  return err_SUCCESS;
}

static error_t client_remove(uint64_t id, client_t** c) {
  *c = (client_t*)id;
  ASSERT_ERR(oe_is_within_enclave(*c, sizeof(client_t)), err_ENCLAVE__GENERAL__CLIENT_REMOVE_FAILED);
  ASSERT_ERR(((uintptr_t)*c) != ((uintptr_t)&g_client_secret), err_ENCLAVE__GENERAL__CLIENT_REMOVE_FAILED);
  ASSERT_ERR(__atomic_compare_exchange_n(&(*c)->canary, &g_client_secret, 0, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQ_REL), err_ENCLAVE__GENERAL__CLIENT_REMOVE_FAILED);
  error_t err = client_setstate(*c, CLIENT_UNUSED, CLIENT_DELETE);
  if (err != err_SUCCESS) {
    __atomic_store_n(&(*c)->canary, g_client_secret, __ATOMIC_RELEASE);
    return err;
  }
  TEST_LOG("client_remove(%p)", *c);
  return err_SUCCESS;
}

int enclave_init(size_t available_memory, double load_factor, size_t num_shards, size_t stash_overflow_size)
{
#ifdef INSECURE
  ENC_LOG_ERROR("!!! WARNING: this enclave is insecure and is intended only for testing purposes");
#endif

  error_t err = err_SUCCESS;
  ASSERT_ERR(load_factor >= 1.0 && load_factor <= 3.0, err_ENCLAVE__GENERAL__INVALID_LOAD_FACTOR);
  ASSERT_ERR(OE_OK == oe_attester_initialize(), err_ENCLAVE__GENERAL__OE_ATTESTER_INITIALIZE);
  // oe_attester_initialize() is basic, and it will succeed as long as one attester plugin loads.
  // We want to fail fast if the format that we use is not available, so we check that it is.
  oe_uuid_t ignored;
  ASSERT_ERR(OE_OK == oe_attester_select_format(&sgx_remote_uuid, 1, &ignored), err_ENCLAVE__GENERAL__OE_ATTESTER_SELECT_FORMAT);
  uint8_t init_expected = ENCLAVE_INIT_UNINITIALIZED;
  ASSERT_ERR(__atomic_compare_exchange_n(&g_init_done, &init_expected, ENCLAVE_INIT_INITIALIZING, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQ_REL), err_ENCLAVE__GENERAL__REINIT);
  GOTO_IF_ERROR(err = generate_privkey(sizeof(g_privkey), g_privkey), done);
  ENC_LOG_DEBUG("enclave_init(available_memory=%zu, load_factor=%lf, num_shards=%zu)", available_memory, load_factor, num_shards);

  g_num_shards = num_shards;

  uint8_t client_map_key[FIXEDMAP_HALFSIPHASH_KEYSIZE];
  CHECK(OE_OK == oe_random(client_map_key, FIXEDMAP_HALFSIPHASH_KEYSIZE));

  signal_user_record record;
  uint64_t t = 0;
  uint8_t hash_key[8];
  CHECK(OE_OK == oe_random(hash_key, 8));
  CHECK(OE_OK == oe_random(&g_client_secret, sizeof(g_client_secret)));
  TEST_LOG("sharded_ohtable_create");
  g_table = sharded_ohtable_create_for_available_mem(sizeof record / sizeof t, num_shards, hash_key, available_memory, load_factor, stash_overflow_size, getentropy);
  memset(hash_key, 0, 8);

done:
  init_expected = ENCLAVE_INIT_INITIALIZING;
  uint8_t init_next = (err == err_SUCCESS) ? ENCLAVE_INIT_INITIALIZED : ENCLAVE_INIT_UNINITIALIZED;
  CHECK(__atomic_compare_exchange_n(&g_init_done, &init_expected, init_next, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQ_REL));
  return err;
}

int enclave_stop_shards()
{
  RETURN_IF_ERROR(check_init_complete());
  for (size_t i = 0; i < g_num_shards; ++i)
  {
    ENC_LOG_DEBUG("Stopping shard %zu\n", i);
    sharded_ohtable_stop_shard(g_table, i);
  }
  return err_SUCCESS;
}

typedef struct {
  char* name;
  uint64_t val;
} counter_t;

static void write_stat_value(struct org_signal_cdsi_value_t* pbval, char* name, uint64_t val) {
  pbval->name_p = name;
  pbval->value = val;
}

int enclave_table_statistics(
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size) {
  RETURN_IF_ERROR(check_init_complete());

  const size_t num_fields = 14;
  const size_t max_field_name_len = 32;
  size_t workspace_size = PBUTIL_WORKSPACE_BASE(struct org_signal_cdsi_table_statistics_t)
    + g_num_shards * (sizeof(struct org_signal_cdsi_shard_statistics_t)
           + num_fields*( sizeof(struct org_signal_cdsi_value_t) + max_field_name_len));
  uint8_t *workspace;
  RETURN_IF_ERROR(MALLOCZ_SIZE(workspace, workspace_size));
  error_t err = err_SUCCESS;

  ohtable_statistics** stats = sharded_ohtable_report_statistics(g_table);

  struct org_signal_cdsi_table_statistics_t *rsp = org_signal_cdsi_table_statistics_new(workspace, workspace_size);
  if (rsp == NULL) {
    err = err_ENCLAVE__TABLE_STATISTICS__RESPONSE_PB_NEW;
    goto finish;
  }
  int allocated = org_signal_cdsi_table_statistics_shard_statistics_alloc(rsp, g_num_shards);
  if (allocated < 0) {
    err = err_ENCLAVE__TABLE_STATISTICS__RESPONSE_PB_ALLOC_SHARDS;
    goto finish;
  }

  for(size_t i = 0; i < g_num_shards; ++i) {
    int alloc_vals = org_signal_cdsi_shard_statistics_values_alloc(rsp->shard_statistics.items_p + i, num_fields);

    if (alloc_vals < 0) {
      err = err_ENCLAVE__TABLE_STATISTICS__RESPONSE_PB_ALLOC_SHARDS;
      goto finish;
    }
    TEST_LOG("num_items: %zu max_overflow: %zu max_trace: %zu mean_overflow: %lf",
            stats[i]->num_items, stats[i]->max_stash_overflow_count, stats[i]->max_trace_length, ((double)stats[i]->sum_stash_overflow_count)/stats[i]->oram_access_count);

    counter_t counters[] = {
      { .name = "max_trace_length", .val = stats[i]->max_trace_length},
      { .name = "total_displacement", .val = stats[i]->total_displacement},
      { .name = "num_items", .val = stats[i]->num_items},
      { .name = "capacity", .val = stats[i]->capacity},
      { .name = "oram_recursion_depth", .val = stats[i]->oram_recursion_depth},
      { .name = "oram_access_count", .val = stats[i]->oram_access_count},
      { .name = "stash_overflow_count", .val = stats[i]->stash_overflow_count},
      { .name = "max_stash_overflow_count", .val = stats[i]->max_stash_overflow_count},
      { .name = "sum_stash_overflow_count", .val = stats[i]->sum_stash_overflow_count},
      { .name = "posmap_stash_overflow_count", .val = stats[i]->posmap_stash_overflow_count},
      { .name = "posmap_max_stash_overflow_count", .val = stats[i]->posmap_max_stash_overflow_count},
      { .name = "posmap_stash_overflow_count", .val = stats[i]->posmap_stash_overflow_count},
      // The next two fields are doubles - exponential moving averages of stash overflow size with a half-life of 10,000.
      // We want the double precision for computation, but really only care about the first few digits for reporting.
      // Instead of creating floating point values for the protobuf, multiply by 10K and store as an integer.
      { .name = "stash_overflow_ema10k", .val = stats[i]->stash_overflow_ema10k * 10000},
      { .name = "posmap_stash_overflow_ema10k", .val = stats[i]->posmap_stash_overflow_ema10k * 10000},
      { .name = NULL }
    };

    for(size_t j = 0; counters[j].name != NULL; ++j) {
      write_stat_value(rsp->shard_statistics.items_p[i].values.items_p + j, counters[j].name, counters[j].val);
    }
  }

  int size = org_signal_cdsi_table_statistics_encode(rsp, out, out_size);
  if(size < 0) {
    err = err_ENCLAVE__TABLE_STATISTICS__RESPONSE_PB_ENCODE;
    goto finish;
  }
  *actual_out_size = size;
finish:
  if(stats) {
    for(size_t i = 0; i < g_num_shards; ++i) {
      ohtable_statistics_destroy(stats[i]);
    }
    free(stats);
  }
  free(workspace);
  return err;
}

int enclave_load_pb(
    size_t len,
    unsigned char *load_request_pb)
{
  RETURN_IF_ERROR(check_init_complete());
  // TEST_LOG("load proto");
  size_t workspace_size = len + PBUTIL_WORKSPACE_BASE(struct org_signal_cdsi_enclave_load_t);
  uint8_t *workspace;
  RETURN_IF_ERROR(MALLOCZ_SIZE(workspace, workspace_size));
  struct org_signal_cdsi_enclave_load_t *load_req = org_signal_cdsi_enclave_load_new(workspace, workspace_size);
  if (load_req == 0)
  {
    return err_ENCLAVE__LOADPB__REQUEST_PB_NEW;
  }
  int size = org_signal_cdsi_enclave_load_decode(load_req, load_request_pb, len);
  if (size < 0)
  {
    return err_ENCLAVE__LOADPB__REQUEST_PB_DECODE;
  }
  if (load_req->shared_token_secret.size > 0)
  {
    RETURN_IF_ERROR(ratelimit_set_shared_secret(
        load_req->shared_token_secret.size,
        load_req->shared_token_secret.buf_p));
  }
  if (load_req->clear_all)
  {
    // clear the table
    TEST_LOG("sharded_ohtable_clear");
    sharded_ohtable_clear(g_table);
  }

  signal_user_record *incoming_records = (signal_user_record *)load_req->e164_aci_pni_uak_tuples.buf_p;
  size_t num_records = load_req->e164_aci_pni_uak_tuples.size / (sizeof *incoming_records);
  if (load_req->e164_aci_pni_uak_tuples.size % (sizeof *incoming_records) != 0)
  {
    return err_ENCLAVE__LOADPB__TUPLES_INVALID;
  }

  // TEST_LOG("sharded_ohtable_put_batch:%zu", num_records);
  error_t result = sharded_ohtable_put_batch(g_table, num_records, (u64 *)incoming_records);
  free(workspace);
  return result;
}

static void client_free(client_t* c) {
  CHECK(oe_is_within_enclave(c, sizeof(client_t)));
  // This may rewrite an already-zero value.
  __atomic_store_n(&c->canary, 0, __ATOMIC_RELEASE);
  __atomic_store_n(&c->state, CLIENT_DELETE, __ATOMIC_RELEASE);
  if (c->workspace)
    free(c->workspace);
  if (c->handshake != NULL)
    noise_handshakestate_free(c->handshake);
  if (c->send != NULL)
    noise_cipherstate_free(c->send);
  if (c->recv != NULL)
    noise_cipherstate_free(c->recv);
  free(c);
}

int enclave_close_client(
    uint64_t cli)
{
  RETURN_IF_ERROR(check_init_complete());
  client_t *c;
  RETURN_IF_ERROR(client_remove(cli, &c));
  TEST_LOG("close: %p", c);
  client_free(c);
  return err_SUCCESS;
}

int enclave_new_client(
    uint64_t *cli,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size)
{
  RETURN_IF_ERROR(check_init_complete());
  client_t *c;
  RETURN_IF_ERROR(MALLOCZ(c));
  TEST_LOG("noise_handshakestate_new_by_name '%s'", NOISE_PROTOCOL_DEFINITION);
  error_t err = err_SUCCESS;
  GOTO_IF_ERROR(err = noise_errort(
      err_NOISE__HANDSHAKESTATE__NEW__,
      noise_handshakestate_new_by_name(
          &c->handshake, NOISE_PROTOCOL_DEFINITION, NOISE_ROLE_RESPONDER)),
      free_client);
  NoiseDHState *local_keypair = noise_handshakestate_get_local_keypair_dh(c->handshake);

  TEST_LOG("client: %p", c);
  CHECK(0 == pthread_rwlock_rdlock(&g_handshakestart_mu));
  GOTO_IF_ERROR(err = validate_privkey(sizeof(g_privkey), g_privkey), unlock_handshake);
  if (out_size < g_handshakestart_size)
  {
    err = err_ENCLAVE__NEWCLIENT__EREPORT_TOO_LARGE;
    goto unlock_handshake;
  }
  else
  {
    GOTO_IF_ERROR(err = noise_errort(err_NOISE__DHSTATE__SET_KEYPAIR_PRIVATE__, noise_dhstate_set_keypair_private(local_keypair, g_privkey, sizeof(g_privkey))), unlock_handshake);
    *actual_out_size = g_handshakestart_size;
    memcpy(out, g_handshakestart, g_handshakestart_size);
  }
  GOTO_IF_ERROR(err = client_add(c, cli), unlock_handshake);
  TEST_LOG("client out size: %zu", *actual_out_size);

  client_setstate(c, CLIENT_INUSE, CLIENT_UNUSED);
unlock_handshake:
  pthread_rwlock_unlock(&g_handshakestart_mu);
free_client:
  if (err != err_SUCCESS) client_free(c);
  return err;
}

int enclave_retry_response(
    uint64_t cli,
    uint32_t retry_after_secs,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size)
{
  RETURN_IF_ERROR(check_init_complete());
  client_t *c;
  RETURN_IF_ERROR(client_get(cli, &c));
  error_t err = err_SUCCESS;
  if (c->send == NULL || c->recv == NULL)
  {
    err = err_ENCLAVE__GENERAL__CLIENT_STATE;
    goto client_unlock;
  }

  TEST_LOG("enclave_retry_response(): proto response encoding into %zu bytes", out_size);
  unsigned char workspace[128];
  struct org_signal_cdsi_client_response_t *rsp = org_signal_cdsi_client_response_new(workspace, sizeof(workspace));
  GOTO_IF_ERROR(err = (rsp != NULL ? err_SUCCESS : err_ENCLAVE__RETRYRESPONSE__RESPONSE_PB_NEW), client_unlock);
  rsp->retry_after_secs = retry_after_secs;

  // Encrypt response containing a retry-after
  uint8_t *plaintext_buf;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(plaintext_buf, out_size), client_unlock);
  *actual_out_size = out_size;
  int rsp_size = org_signal_cdsi_client_response_encode(rsp, plaintext_buf, out_size);
  if (rsp_size < 0)
  {
    TEST_LOG("enclave_retry_response(): org_signal_cdsi_client_response_encode failure: %d", rsp_size);
    err = err_ENCLAVE__RETRYRESPONSE__RESPONSE_PB_ENCODE;
  }
  else
  {
    err = noise_encrypt_message(c->send, plaintext_buf, rsp_size, out, actual_out_size);
  }
  free(plaintext_buf);
client_unlock:
  client_setstate(c, CLIENT_INUSE, CLIENT_UNUSED);
  return err;
}

int enclave_rate_limit(
    uint64_t cli,
    size_t in_size,
    unsigned char *in,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size,
    size_t old_token_hash_size,
    unsigned char *old_token_hash,
    size_t *actual_old_token_hash_size,
    size_t new_token_hash_size,
    unsigned char *new_token_hash,
    size_t *actual_new_token_hash_size,
    size_t *computed_request_size)
{
  RETURN_IF_ERROR(check_init_complete());
  error_t err = err_SUCCESS;
  client_t *c;
  RETURN_IF_ERROR(client_get(cli, &c));
  if (c->req != NULL || c->send == NULL || c->recv == NULL || c->workspace != NULL)
  {
    err = err_ENCLAVE__GENERAL__CLIENT_STATE;
    goto client_unlock;
  }
  GOTO_IF_ERROR(err = (old_token_hash_size >= RATELIMIT_TOKENHASH_SIZE ? err_SUCCESS : err_ENCLAVE__RATELIMIT__BAD_OUTPUT), client_unlock);
  GOTO_IF_ERROR(err = (new_token_hash_size >= RATELIMIT_TOKENHASH_SIZE ? err_SUCCESS : err_ENCLAVE__RATELIMIT__BAD_OUTPUT), client_unlock);
  // Decrypt initial client request.
  TEST_LOG("noise_decrypt_message_inplace size=%zu", in_size);
  GOTO_IF_ERROR(err = noise_decrypt_message_inplace(c->recv, in, &in_size), client_unlock);

  // Decode initial client request.
  TEST_LOG("proto request size=%zu", in_size);
  size_t workspace_size = in_size  + PBUTIL_WORKSPACE_BASE(struct org_signal_cdsi_client_request_t);
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(c->workspace, workspace_size), client_unlock);
  struct org_signal_cdsi_client_request_t *req = org_signal_cdsi_client_request_new(c->workspace, workspace_size);
  if (req == 0)
    return err_ENCLAVE__RATELIMIT__REQUEST_PB_NEW;
  int size = org_signal_cdsi_client_request_decode(req, in, in_size);
  if (size < 0)
    return err_ENCLAVE__RATELIMIT__REQUEST_PB_DECODE;

  // Check input preconditions using constant-time operations. We require that:
  //
  // - aci_uak_pairs has a length that's a multiple of 32 bytes (the ACI is a UUID, which is 16 bytes; UAKs are, by
  //   definition, 16 bytes)
  // - prev_e164s has a length that's a multiple of 8 bytes (each e164 is 8 bytes)
  // - new_e164s has a length that's a multiple of 8 bytes
  // - At least one of new_e164s or prev_e164s has a non-zero length
  if (((req->aci_uak_pairs.size & (4 * sizeof(uint64_t) - 1)) != 0) |
      ((req->prev_e164s.size & (sizeof(uint64_t) - 1)) != 0) |
      ((req->new_e164s.size & (sizeof(uint64_t) - 1)) != 0) |
      req->new_e164s.size + req->prev_e164s.size == 0)
  {
    TEST_LOG("Invalid argument to enclave_rate_limit. Invalid array size");
    return err_ENCLAVE__RATELIMIT__REQUEST_PB_INVALID;
  }
  c->req = req;

  // Check old token.  If valid, set computed output size.
  TEST_LOG("ratelimit_validate_received_rate_limit_token");
  RETURN_IF_ERROR(ratelimit_validate_received_rate_limit_token(c->req));

  // divide the size in bytes by 8 to get the size in uint64_ts. Use shift for
  // fast constant time division by power of 2.
  *computed_request_size = req->new_e164s.size >> 3;

  // Compute new token
  uint8_t new_token[RATELIMIT_TOKEN_SIZE] = {RATELIMIT_VERSION};
  GOTO_IF_ERROR(err = (OE_OK == oe_random(new_token+RATELIMIT_VERSION_SIZE, RATELIMIT_ENTROPY_SIZE) ? err_SUCCESS : err_ENCLAVE__GENERAL__ENTROPY), client_unlock);
  TEST_LOG("ratelimit_compute_token");
  GOTO_IF_ERROR(err = ratelimit_compute_token(
      &req->prev_e164s,
      &req->new_e164s,
      new_token), client_unlock);

  // Set token hashes
  TEST_LOG("ratelimit_token_hash 1");
  if (req->token.size == 0)
  {
    // If no old token, then "hash" is empty string
    *actual_old_token_hash_size = 0;
  }
  else
  {
    ratelimit_token_hash(req->token.buf_p, old_token_hash);
    *actual_old_token_hash_size = RATELIMIT_TOKENHASH_SIZE;
  }
  TEST_LOG("ratelimit_token_hash 2");
  ratelimit_token_hash(new_token, new_token_hash);
  *actual_new_token_hash_size = RATELIMIT_TOKENHASH_SIZE;

  // Generate response that provides new_token to client.
  TEST_LOG("proto response encoding into %zu bytes", out_size);
  unsigned char workspace[128];
  struct org_signal_cdsi_client_response_t *rsp =
      org_signal_cdsi_client_response_new(workspace, sizeof(workspace));
  GOTO_IF_ERROR(err = (rsp != NULL ? err_SUCCESS : err_ENCLAVE__RATELIMIT__RESPONSE_PB_NEW), client_unlock);
  rsp->token.buf_p = new_token;
  rsp->token.size = sizeof(new_token);

  uint8_t *plaintext_buf;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(plaintext_buf, out_size), client_unlock);
  *actual_out_size = out_size;
  int rsp_size = org_signal_cdsi_client_response_encode(rsp, plaintext_buf, out_size);
  if (rsp_size < 0)
  {
    ENC_LOG_ERROR("org_signal_cdsi_client_response_encode failure: %d", rsp_size);
    err = err_ENCLAVE__RATELIMIT__RESPONSE_PB_ENCODE;
  }
  else
  {
    err = noise_encrypt_message(c->send, plaintext_buf, rsp_size, out, actual_out_size);
  }
  free(plaintext_buf);

client_unlock:
  client_setstate(c, CLIENT_INUSE, CLIENT_UNUSED);
  return err;
}

int enclave_run(
    uint64_t cli,
    uint32_t permits,
    size_t in_size,
    unsigned char *in,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size)
{
  RETURN_IF_ERROR(check_init_complete());
  client_t *c;
  RETURN_IF_ERROR(client_get(cli, &c));
  if (c->req == NULL || c->send == NULL || c->recv == NULL)
  {
    return err_ENCLAVE__GENERAL__CLIENT_STATE;
  }
  error_t err = err_SUCCESS;

  // Decrypt incoming buffer.
  TEST_LOG("noise_decrypt_message_inplace size=%zu", in_size);
  GOTO_IF_ERROR(err = noise_decrypt_message_inplace(c->recv, in, &in_size), client_unlock);

  // Verify that the client has ack'd the token.
  unsigned char req_workspace[128];
  TEST_LOG("proto request size=%zu", in_size);
  struct org_signal_cdsi_client_request_t *req = org_signal_cdsi_client_request_new(req_workspace, sizeof(req_workspace));
  GOTO_IF_ERROR(err = (req == 0 ? err_ENCLAVE__RUN__REQUEST_PB_NEW : err_SUCCESS), client_unlock);
  int size = org_signal_cdsi_client_request_decode(req, in, in_size);
  GOTO_IF_ERROR(err = (size < 0 ? err_ENCLAVE__RUN__REQUEST_PB_DECODE : err_SUCCESS), client_unlock);
  GOTO_IF_ERROR(err = (!req->token_ack ? err_ENCLAVE__RUN__NO_TOKEN_ACK : err_SUCCESS), client_unlock);

  // create and populate the aci_uak map
  size_t size_e164s = c->req->prev_e164s.size + c->req->new_e164s.size;
  size_t num_e164s = size_e164s >> 3; // divide by sizeof(uint64_t) in constant time
  size_t num_aci_uak_pairs = c->req->aci_uak_pairs.size >> 5; // divide by  (4 * sizeof(uint64_t)) in constant time

  TEST_LOG("aci_uak map");
  fixedset_t *index = NULL;
  unsigned char sipkey[FIXEDMAP_HALFSIPHASH_KEYSIZE];
  oe_random(sipkey, FIXEDMAP_HALFSIPHASH_KEYSIZE);
  GOTO_IF_ERROR(err = fixedset_new(&index, sizeof(aci_uak_pair), sipkey), client_unlock);
  GOTO_IF_ERROR(err = load_aci_uak_pairs(index, num_aci_uak_pairs, (aci_uak_pair *)c->req->aci_uak_pairs.buf_p), free_index);

  // --- Prepare and execute query ---
  // merge e164s
  // The following block has the same effect as two memcpy calls:
  //
  //   memcpy(e164s, c->req->prev_e164s.buf_p, c->req->prev_e164s.size);
  //   memcpy(e164s + c->req->prev_e164s.size, c->req->new_e164s.buf_p, c->req->new_e164s.size);
  //
  // However it performs the operations with oblivious writes to prevent leaking the size of the
  // `new_e164s` array.
  uint8_t *e164s;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(e164s, size_e164s), free_index);
  size_t prev_e164s_size = c->req->prev_e164s.size / sizeof(uint64_t);
  uint64_t* new_e164s = c->req->new_e164s.buf_p == NULL
      // We must be able to reference *(new_e164s+i-prev_e164s_size) for all i in [0,num_e164s).
      // If buf_p is null, then all e164s are within prev_e164s, so pointing to the END of
      // that buffer maintains this invariant.
      ? (uint64_t*)(c->req->prev_e164s.buf_p + c->req->prev_e164s.size)
      : (uint64_t*)c->req->new_e164s.buf_p;
  uint64_t* prev_e164s = c->req->prev_e164s.buf_p == NULL
      // We must be able to reference (*prev_e164s+i) for all i in [0,num_e164s).
      // If buf_p is null, then all e164s are within new_e164s, so pointing to the BEGINNING
      // of that buffer maintains this invariant.
      ? new_e164s
      : (uint64_t*)c->req->prev_e164s.buf_p;
  for(size_t i = 0; i < num_e164s; ++i) {
    // This assigns the value of the i-th e164 to be either (1) the i-th prev_e164 if
    // i < prev_e164s.size or(2) the (i - prev_e164s)-th n_e164 otherwise.
    //
    // The indexes used look as if they should cause memory violations, but this is safe because
    // (1) all memory is part of the client workspace (c->workspace) allocated in `enclave_rate_limit` and
    // (2) the new_e164s are allocated after the prev_e164s in that array.
    ((uint64_t*)e164s)[i] = U64_TERNARY(i < prev_e164s_size, *(prev_e164s + i), *(new_e164s + i - prev_e164s_size));
  }

  signal_user_record *user_records;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(user_records, num_e164s * sizeof(*user_records)), free_e164s);
  TEST_LOG("sharded_ohtable_get_batch size=%zu", sizeof(*user_records) * num_e164s);
  GOTO_IF_ERROR(err = sharded_ohtable_get_batch(g_table, num_e164s, (u64 *)e164s, (u64 *)user_records), free_e164s);

  // process response
  e164_pni_aci_triple *out_triples;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(out_triples, num_e164s * sizeof(*out_triples)), free_user_records);
  GOTO_IF_ERROR(
      err = create_e164_pni_aci_triples(
          index,
          num_e164s,
          (uint64_t *)e164s,
          user_records,
          out_triples),
      free_user_records);

  // Create space for response.
  TEST_LOG("proto response");
  size_t response_workspace_size = PBUTIL_WORKSPACE_BASE(struct org_signal_cdsi_client_response_t);
  uint8_t *response_workspace;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(response_workspace, response_workspace_size), free_out_triples);
  struct org_signal_cdsi_client_response_t *rsp = org_signal_cdsi_client_response_new(response_workspace, response_workspace_size);

  if (rsp == NULL)
  {
    err = err_ENCLAVE__RUN__RESPONSE_PB_NEW;
    goto free_response_workspace;
  }
  // Set response fields and encode.
  rsp->e164_pni_aci_triples.buf_p = (uint8_t *)out_triples;
  rsp->e164_pni_aci_triples.size = num_e164s * sizeof(*out_triples);
  rsp->token.buf_p = 0;
  rsp->token.size = 0;
  rsp->debug_permits_used = permits;

  uint8_t *plaintext_buf;
  GOTO_IF_ERROR(err = MALLOCZ_SIZE(plaintext_buf, out_size), free_response_workspace);

  int resp_size = org_signal_cdsi_client_response_encode(rsp, plaintext_buf, out_size);
  if (resp_size < 0)
  {
    err = err_ENCLAVE__RUN__RESPONSE_PB_ENCODE;
    goto free_plaintext_buf;
  }

  *actual_out_size = out_size; // this will be overwritten by noise_encrypt_message
  err = noise_encrypt_message(c->send, plaintext_buf, resp_size, out, actual_out_size);

free_plaintext_buf:
  free(plaintext_buf);
free_response_workspace:
  free(response_workspace);
free_out_triples:
  free(out_triples);
free_user_records:
  free(user_records);
free_e164s:
  free(e164s);
free_index:
  fixedset_free(index);
client_unlock:
  client_setstate(c, CLIENT_INUSE, CLIENT_UNUSED);
  return err;
}

int enclave_run_shard(size_t shard_id)
{
  RETURN_IF_ERROR(check_init_complete());
  sharded_ohtable_run_shard(g_table, shard_id);
  return err_SUCCESS;
}

static error_t wrap_in_ereport(
    uint8_t* pubkey,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size)
{
  if (out_size <= NOISE_KEY_SIZE) {
    return err_ENCLAVE__GENERAL__EREPORT_SIZE;
  }
  uint8_t workspace[32 << 10];
  error_t err = err_SUCCESS;

  struct org_signal_cdsi_client_handshake_start_t *rsp = org_signal_cdsi_client_handshake_start_new(workspace, sizeof(workspace));
  ASSERT_ERR(rsp != NULL, err_ENCLAVE__GENERAL__EREPORT_PB_NEW);
  rsp->test_only_pubkey.size = NOISE_KEY_SIZE;
  rsp->test_only_pubkey.buf_p = pubkey;

  char custom_claim_name[] = "pk";
  oe_claim_t custom_claims[] = {{.name = custom_claim_name, .value = pubkey, .value_size = NOISE_KEY_SIZE }};
  uint8_t* custom_claims_buffer = NULL;
  size_t custom_claims_buffer_size = 0;
  ASSERT_ERR(
      OE_OK == oe_serialize_custom_claims(
          custom_claims,
          1,
          &custom_claims_buffer,
          &custom_claims_buffer_size),
      err_ENCLAVE__GENERAL__OE_SERIALIZE_CUSTOM_CLAIMS);

  uint8_t * evidence_buffer = NULL;
  size_t evidence_buffer_size = 0;
  uint8_t *endorsements_buffer = NULL;
  size_t endorsements_buffer_size = 0;
  oe_result_t result = oe_get_evidence(
      &sgx_remote_uuid, 0,
      custom_claims_buffer, custom_claims_buffer_size,
      NULL, 0,
      &evidence_buffer, &evidence_buffer_size,
      &endorsements_buffer, &endorsements_buffer_size);
  if (result != OE_OK) {
    TEST_LOG("get evidence error: %s", oe_result_str(result));
#if defined IS_TEST || defined INSECURE
    ENC_LOG_WARN("Continuing without evidence, endorsements");
#else
    err = err_ENCLAVE__GENERAL__GET_EVIDENCE;
    goto free_custom_claim;
#endif
  } else {
    rsp->evidence.buf_p = evidence_buffer;
    rsp->evidence.size = evidence_buffer_size;
    rsp->endorsement.buf_p = endorsements_buffer;
    rsp->endorsement.size = endorsements_buffer_size;
  }

  int rsp_size = org_signal_cdsi_client_handshake_start_encode(rsp, out, out_size);
  if (rsp_size < 0) {
    err = err_ENCLAVE__GENERAL__EREPORT_PB_ENCODE;
    goto free_evidence;
  }

  *actual_out_size = rsp_size;

free_evidence:
  if (endorsements_buffer) oe_free_endorsements(endorsements_buffer);
  if (evidence_buffer) oe_free_evidence(evidence_buffer);
#if !defined IS_TEST && !defined INSECURE
free_custom_claim:
#endif
  oe_free_serialized_custom_claims(custom_claims_buffer);
  return err;
}

int enclave_handshake(
    uint64_t cli,
    size_t in_size,
    unsigned char *in,
    size_t out_size,
    unsigned char *out,
    size_t *actual_out_size)
{
  RETURN_IF_ERROR(check_init_complete());
  client_t *c;
  RETURN_IF_ERROR(client_get(cli, &c));
  error_t err = err_SUCCESS;
  if (c->handshake == NULL) {
    err = err_ENCLAVE__GENERAL__CLIENT_STATE;
    goto client_unlock;
  }

  TEST_LOG("noise_handshakestate_start");
  GOTO_IF_ERROR(err = noise_errort(err_NOISE__HANDSHAKESTATE__START__, noise_handshakestate_start(c->handshake)), client_unlock);

  // Read in and use the client's message as a NoiseNN handshake request.
  NoiseBuffer buf_in;
  noise_buffer_set_input(buf_in, in, in_size);
  TEST_LOG("noise_handshakestate_read_message");
  if (NOISE_ACTION_READ_MESSAGE != noise_handshakestate_get_action(c->handshake)) {
    err = err_ENCLAVE__HANDSHAKE__INVALID_STATE;
    goto client_unlock;
  }
  GOTO_IF_ERROR(err = noise_errort(err_NOISE__HANDSHAKESTATE__READ__, noise_handshakestate_read_message(c->handshake, &buf_in, NULL)), client_unlock);

  // Generate the handshake response.
  unsigned char handshake_msg[64];
  NoiseBuffer buf_out;
  noise_buffer_set_output(buf_out, handshake_msg, sizeof(handshake_msg));
  TEST_LOG("noise_handshakestate_write_message");
  if (NOISE_ACTION_WRITE_MESSAGE != noise_handshakestate_get_action(c->handshake)) {
    err = err_ENCLAVE__HANDSHAKE__INVALID_STATE;
    goto client_unlock;
  }
  GOTO_IF_ERROR(err = noise_errort(err_NOISE__HANDSHAKESTATE__WRITE__, noise_handshakestate_write_message(c->handshake, &buf_out, NULL)), client_unlock);
  TEST_LOG("handshake: size=%zu", buf_out.size);

  memcpy(out, handshake_msg, buf_out.size);
  *actual_out_size = buf_out.size;

  // Split the handshake.
  TEST_LOG("noise_handshakestate_split");
  if (NOISE_ACTION_SPLIT != noise_handshakestate_get_action(c->handshake)) {
    err = err_ENCLAVE__HANDSHAKE__INVALID_STATE;
    goto client_unlock;
  }
  GOTO_IF_ERROR(err = noise_errort(err_NOISE__HANDSHAKESTATE__SPLIT__, noise_handshakestate_split(c->handshake, &c->send, &c->recv)), client_unlock);
  noise_handshakestate_free(c->handshake);
  c->handshake = NULL;
  TEST_LOG("success");
client_unlock:
  client_setstate(c, CLIENT_INUSE, CLIENT_UNUSED);
  return err;
}

int enclave_attest()
{
  RETURN_IF_ERROR(check_init_complete());
  uint8_t new_key[NOISE_KEY_SIZE];
  RETURN_IF_ERROR(generate_privkey(sizeof(new_key), new_key));
  uint8_t pubkey[NOISE_KEY_SIZE];
  RETURN_IF_ERROR(noise_pubkey_from_privkey(new_key, pubkey));
  if (sizeof(g_handshakestart) < sizeof(pubkey))
  {
    ENC_LOG_ERROR("Bad ereport size (%zu) when renewing attestation", sizeof(g_handshakestart));
    return err_ENCLAVE__GENERAL__EREPORT_SIZE;
  }
  uint8_t* scratch;
  size_t scratch_size = sizeof(g_handshakestart);
  RETURN_IF_ERROR(MALLOCZ_SIZE(scratch, scratch_size));
  error_t err = err_SUCCESS;
  GOTO_IF_ERROR(
      err = wrap_in_ereport(pubkey, scratch_size, scratch, &scratch_size),
      free_scratch);
  pthread_rwlock_wrlock(&g_handshakestart_mu);
  memcpy(g_privkey, new_key, sizeof(new_key));
  memcpy(g_handshakestart, scratch, scratch_size);
  g_handshakestart_size = scratch_size;
  pthread_rwlock_unlock(&g_handshakestart_mu);
  TEST_LOG("got ereport size: %zu\n", g_handshakestart_size);
free_scratch:
  free(scratch);
  return err;
}
